#include "Genome_Index_TableQ.h"

CGenome_Index_TableQ::CGenome_Index_TableQ(void)
{
    initialization();
}

CGenome_Index_TableQ::~CGenome_Index_TableQ(void)
{
    // this->pgenomeNTInBits and this->pgenome will be deleted in parent class;
}

int CGenome_Index_TableQ::initialization(void)
{
    this->bExcludeAmbiguous = true;
    return(0);
}

bool CGenome_Index_TableQ::getSeqFromFasta(const char* genomeListfileName, string refFormat)
{
    getBasename(genomeListfileName, this->caRefName);
    if (fileExist(genomeListfileName)) {
        this->getGenomeNTdata(genomeListfileName, refFormat);
        TIME_INFO(this->pgenomeNTInBits = new CGenomeInBits(this->pgenomeNT), "Build genome in bits ");
        this->pgenomeNT->freeChromosomeSpace();
        return(true);
    } else {
        cout << "File " << genomeListfileName << " not found!" << endl;
        return(false);
    }
}

bool CGenome_Index_TableQ::getSeqFromDS(CGenomeNTdata* pgenomeNT)
{
    this->pgenomeNT = pgenomeNT;
    myStrCpy(this->caRefName, pgenomeNT->refName, FILENAME_MAX);
    TIME_INFO(this->pgenomeNTInBits = new CGenomeInBits(this->pgenomeNT), "Build genome in bits ");
    this->pgenomeNT->freeChromosomeSpace();
    return(true);
}

// (1) Query each seed pattern for hit. (2) Check each hit for valid alignment and put in the palignmentsQ
// (3) Return the minDiff for alignment
pair<CIndex_Type*, CIndex_Type*> CGenome_Index_TableQ::queryKmer\
(CReadInBits window, unsigned int shift) const
{
    unsigned int BIN_SIZE_CHECK_THRESHOLD = 1; // If there are only this number in bin, check

    window = window.getSuffixStr(shift);

    unsigned int uiHashValue = this->getHashValue(window);
    unsigned int uiSeedKeyL = this->getSeedKey(window);
    unsigned int uiSeedKeyU;
    if (bEXTEND_SEED) {
        uiSeedKeyU = this->getSeedKeyUpperBound(window, shift);
    } else {
        uiSeedKeyU = uiSeedKeyL;
    }
    // cout << uiHashValue << ',' << uiSeedKey << endl;

    unsigned int bucketStart = this->pHashIndexTable->aiIndexTable[uiHashValue];
    unsigned int nextBucketStart = this->pHashIndexTable->aiIndexTable[uiHashValue + 1];

    if (bucketStart > nextBucketStart) ERR; // Error Case
    if (bucketStart >= nextBucketStart) { // empty bucket
        return(pair<CIndex_Type*, CIndex_Type*>(NULL, NULL)); //No hit
    }

    CIndex_Type* hitStart = &this->pIndexTable[bucketStart];
    CIndex_Type* hitEndPlus1 = &this->pIndexTable[nextBucketStart];

    // If the bin is large enough, use binary search to find the correct hituiSeedKey
    if (bucketStart + BIN_SIZE_CHECK_THRESHOLD  < nextBucketStart) {
        hitStart    = lower_bound(hitStart, hitEndPlus1, uiSeedKeyL, CcompareFunctor4LowerBound(this));
        // TODO fill the suffix digit for the uiSeedKey
        hitEndPlus1 = upper_bound(hitStart, hitEndPlus1, uiSeedKeyU, CcompareFunctor4UpperBound(this));
    } // else assume everything in the bucket is a hit without equal_range check
    pair<CIndex_Type*, CIndex_Type*> returnV(hitStart, hitEndPlus1);
    return(returnV);
}

unsigned int CGenome_Index_TableQ::getSeedKeyUpperBound(CReadInBits window, unsigned int shift) const
{
    //WORD_SIZE padding = 0xffffffffffffffff;
    WORD_SIZE padding = -1;
    if (this->bMapReadInColors) {
        padding <<= (this->uiRead_Length - 1 - shift);

    } else {
        padding <<= (this->uiRead_Length - shift);
    }
    window.UpperBits |= padding;
    window.LowerBits |= padding;

    return(this->getSeedKey(window));
}

// Query a read in bases (Illumina) for hit and check uiDiff and put the result into the given Queue
unsigned int CGenome_Index_TableQ::queryReadBases(CReadInBits readInBases, CAlignmentsQ& alignmentsQ, bool bClearQ, bool bForward) const
{
    if (bClearQ) {
        alignmentsQ.clearHits();
    }
    if (!bForward) { //Query reverse compliment read for alignment
        readInBases = reverseCompliment(this->uiRead_Length, readInBases);
    }
    for (unsigned int shift = 0; shift <= this->uiNoOfShift; shift++) {
        // If the best match is exact matched, no need to go to the next shift, for all exact matches must be found previously.
        if ((alignmentsQ.MinDiff == 0) && (shift > 0) && (alignmentsQ.load > 0)) {
            break;
        }
        pair<CIndex_Type*, CIndex_Type*> hits = queryKmer(readInBases, shift);
        if (hits.first != NULL) {
            for (CIndex_Type* it = hits.first; it < hits.second; it++) {
                if (*it >= shift) {
                    unsigned int alignStart = *it - shift;
                    if (isMasked(alignStart)) {
                        continue;
                    } else {
                        CReadInBits ref = this->pgenomeNTInBits->getSubstringInBits(alignStart);
                        unsigned int uiDiff = bitsStrNCompare(ref, readInBases, this->uiRead_Length);
                        // The flag in alignmentsQ decide whether all alignments within uiDiff or
                        // only the alignments with minimum Diff are queue.
                        if (uiDiff <= this->uiSubDiffThreshold) {
                            alignmentsQ.pushHits(alignStart, (unsigned short)uiDiff, bForward);
                        }

                        if (this->bExcludeAmbiguous) {
                            // short cut to exclude ambiguous reads
                            bool bMap2Repeat = this->pbaRepeatRepresentativeFlag->b(alignStart);
                            if (bMap2Repeat) {
                                alignmentsQ.AmbiguousFlag = true;
                            }
                            if (alignmentsQ.MinDiff == 0 && (alignmentsQ.load >= 2 || bMap2Repeat)) {
                                alignmentsQ.setForwardLoad(bForward);
                                return(0);
                            }
                        }
                    }
                }
            }
            if (alignmentsQ.MinDiff == 0) {
                // short cut. Output no more than iMaxCapacity exaxt alignment
                if (alignmentsQ.load >= alignmentsQ.iMaxCapacity - 1) {
                    alignmentsQ.setForwardLoad(bForward);
                    return(0);
                }
                // short cut. All exact matches will be found after first shift
                if (this->bExcludeAmbiguous && alignmentsQ.load >= 2) {
                    alignmentsQ.setForwardLoad(bForward);
                    return(1);
                }
            }
        }
    }
    /* DEBUG
    if(palignmentsQ.MinDiff <= this->uiSubDiffThreshold) {
        CReadInBits ref = this->pgenomeNTInBits->getSubstringInBits(alignmentsQ.aiHitIndex[0]);
        printBitsStrCompare(ref, readInBases, "Good Alignments!!");
    }*/

    // The records before are for the forward direction if this is a forward query
    alignmentsQ.setForwardLoad(bForward);
    return(alignmentsQ.MinDiff);
}

// Given alignments in alignmentsQ, check reads can be also aligned in the extended position
unsigned int CGenome_Index_TableQ::extendAlignment\
(CAlignmentsQ& alignmentsQ, CReadInBits extendedReadHalf) const
{
    unsigned int i;
    int minDiff = this->uiRead_Length;
    bool extendForward = true;
    for (i = 0; i < alignmentsQ.ForwardAlignmentLoad; i++) {
        unsigned int alignStart = alignmentsQ.aiHitIndex[i] + this->uiRead_Length;
        if (alignStart + this->uiRead_Length >= this->pgenomeNT->iGenomeSize || this->isMasked(alignStart)) {
            alignmentsQ.asdiff[i] += (unsigned short)this->uiRead_Length; // no space to align
        } else {
            CReadInBits ref = this->pgenomeNTInBits->getSubstringInBits(alignStart, this->uiRead_Length);
            alignmentsQ.asdiff[i] += (unsigned short)bitsStrNCompare(ref, extendedReadHalf, this->uiRead_Length);
            if (alignmentsQ.asdiff[i] < minDiff) {
                minDiff = alignmentsQ.asdiff[i];
            }
        }
    }
    extendForward = false;
    CReadInBits reverseComplimentRead = reverseCompliment(this->uiRead_Length, extendedReadHalf);
    for (; i < alignmentsQ.load; i++) { // reverse read
        if ( alignmentsQ.aiHitIndex[i] <  this->uiRead_Length) {
            alignmentsQ.asdiff[i] += (unsigned short)this->uiRead_Length; // no space to align
        } else {
            unsigned int alignStart = alignmentsQ.aiHitIndex[i] - this->uiRead_Length;
            if ( this->isMasked(alignStart) ) {
                alignmentsQ.asdiff[i] += (unsigned short)this->uiRead_Length; // no space to align
            } else {
                CReadInBits ref = this->pgenomeNTInBits->getSubstringInBits(alignStart, this->uiRead_Length);
                alignmentsQ.asdiff[i] += (unsigned short)bitsStrNCompare(ref, reverseComplimentRead, this->uiRead_Length);
                alignmentsQ.aiHitIndex[i] = alignStart;
                if (alignmentsQ.asdiff[i] < minDiff) {
                    minDiff = alignmentsQ.asdiff[i];
                }
            }
        }
    }
    alignmentsQ.MinDiff = minDiff;
    // The alignments may exceed the threshold so it should be checked outside the function.
    return(alignmentsQ.MinDiff);
}

// Query a read in colors (SOLiD) for hit and check uiDiff and put the result into the given Queue
unsigned int CGenome_Index_TableQ::queryReadColors(CReadInBits readInColors, CAlignmentsQ& alignmentsQ, bool bClearQ, bool bForward) const
{
    if (bClearQ) {
        alignmentsQ.clearHits();
    }

    CReadInBits pureColors = readInColors.getSuffixStr(1); // query colors
    if (!bForward) { // Query alignment with reversed color. (Reverse compliment reads has color in reversed direction)
        pureColors = reversePureColors(pureColors, this->uiRead_Length - 1);
    }

    for (unsigned int shift = 0; shift <= this->uiNoOfShift; shift++) {
        // If the best alignment is exact matched, no need to go to the next shift, for all exact matches must be found previously.
        if ((alignmentsQ.MinDiff == 0) && (shift > 0) && (alignmentsQ.load > 0)) {
            break;
        }

        // printBitsStr(pureColors.getSuffixStr(shift), this->uiSeedLength);// DEBUG
        pair<CIndex_Type*, CIndex_Type*> hits = queryKmer(pureColors, shift);
        if (hits.first != NULL) {  // for each hit, check if it is a good alignment.
            for (CIndex_Type* it = hits.first; it < hits.second; it++) {
                if (*it >= shift) {
                    unsigned int alignStart = *it - shift;
                    if (isMasked(alignStart)) {
                        continue;
                    } else {
                        CReadInBits ref = this->pgenomeNTInBits->getSubstringInBits(alignStart, this->uiRead_Length);
                        unsigned int uiDiff;
                        if (bForward) {
                            CReadInBits refInColors = bases2Colors(ref);
                            uiDiff = bitsStrNCompare(refInColors, readInColors, this->uiRead_Length);
                        } else { //reverse compliment the reference to compare in the read
                            CReadInBits rcRefInColors = bases2Colors(reverseCompliment(this->uiRead_Length, ref));
                            uiDiff = bitsStrNCompare(rcRefInColors, readInColors, this->uiRead_Length);
                        }
                        // To queue all alignment within uiDiff or the alignments with min difference are set in flag of alignmentsQ
                        if (uiDiff <= this->uiSubDiffThreshold) {
                            alignmentsQ.pushHits(alignStart, (unsigned short)uiDiff, bForward);
                        }

                        if (this->bExcludeAmbiguous) {
                            // short cut to exclude ambiguous reads
                            bool bMap2Repeat = this->pbaRepeatRepresentativeFlag->b(alignStart);
                            if (bMap2Repeat) {
                                alignmentsQ.AmbiguousFlag = true;
                            }
                            if (alignmentsQ.MinDiff == 0 && (alignmentsQ.load >= 2 || bMap2Repeat)) {
                                return(0);
                            }
                        }
                    }
                }
            }
            // short cut to exclude low entropy reads that has too many ambiguous mapping
            const unsigned int MAX_NO_OF_ALIGMNET = 100;
            if (alignmentsQ.load >= MAX_NO_OF_ALIGMNET && alignmentsQ.MinDiff <= 1) {
                return(alignmentsQ.MinDiff);
            }

            if (this->bExcludeAmbiguous && alignmentsQ.MinDiff <= 1 && alignmentsQ.load >= 2) {
                return(1);
            }
        }
    }

    if (bForward) { // Query reverse compliment read for alignment
        alignmentsQ.ForwardAlignmentLoad = alignmentsQ.load; // The records before are for the forward direction
    }
    return(alignmentsQ.MinDiff);
}

void introduceMutation(char* kmer, int mutationPattern)
{
    if (mutationPattern == 2 || mutationPattern == 3 || mutationPattern == 4) {
        mutateRead(kmer, mutationPattern);
    } else if (mutationPattern == (int)FULL_SENSITIVE_OPT_TO_ONE_BASE_ONE_COLOR_MIS) {
        mutatePairsOfConsecutiveBases(kmer, 1);
        mutateRead(kmer, 1);
    } else if (mutationPattern == (int)FULL_SENSITIVE_OPT_TO_TWO_BASE_MIS ) {
        mutatePairsOfConsecutiveBases(kmer, 2);
    } else if (mutationPattern == (int)FULL_SENSITIVE_OPT_TO_ONE_BASE_TWO_COLOR_MIS) {
        mutatePairsOfConsecutiveBases(kmer, 1);
        mutateRead(kmer, 2);
    }
}

/*
 * Test function which mutates every sliding windows on the reference genome and check if the all hits are found
 */
bool testGenome_Index_TableQ(CGenome_Index_TableQ* table)
{
    bool pass = true;    
    CAlignmentsQ alignmentsQ('B');

    CchromosomeNTdata* testChr;
    unsigned int testChrId;
    if (table->pgenomeNT->iNo_of_chromosome > 1) {
        testChrId = 1; // If there are more than one chromosome, test the second one.
        testChr = table->pgenomeNT->paChromosomes[1];
    } else {
        testChrId = 0;
        testChr = table->pgenomeNT->paChromosomes[0];
    }    
    unsigned int testGenomeIdStart = table->pgenomeNT->chrIndex2genomelocusID(testChrId, 0);
    unsigned int testGenomeIdEnd   = table->pgenomeNT->chrIndex2genomelocusID(testChrId, 0)
                                     + testChr->iChromosome_size - table->uiRead_Length;       
    for (unsigned int i = testGenomeIdStart; i < testGenomeIdEnd; i++) {
        CReadInBits slideWindows;
        bool goodRead;
        goodRead = table->pgenomeNTInBits->fragACGTKmerInBits(slideWindows, i, table->uiRead_Length);
        if (table->bMapReadInColors) {
            slideWindows = bases2Colors(slideWindows);
        }

        if (goodRead) {
            char read[wordSize + 1];
            char originalRead[wordSize + 1];
            slideWindows.decode(read);

            CReadInBits readInBits(read);
            alignmentsQ.clearHits();
            // (1) Test exact match
            if (table->bMapReadInColors) {
                table->queryReadColors(readInBits, alignmentsQ, true, true);
            } else {
                table->queryReadBases(readInBits, alignmentsQ, true, true);
            }
            if (alignmentsQ.MinDiff != 0) {
                cout << "Miss exact match " << i << endl;
                pass = false;
                PRINT_MASKED_REPERAT_FLAG(table, i);
            }

            // (2) Test mutated reads
            myStrCpy(originalRead, read, wordSize + 1);

            introduceMutation(read, table->chosenSeedId);
            CReadInBits mReadInBits(read);
            alignmentsQ.clearHits();

            if (table->bMapReadInColors) {
                table->queryReadColors(mReadInBits, alignmentsQ, true, true);
            } else {
                table->queryReadBases(mReadInBits, alignmentsQ, true, true);
            }
            if (alignmentsQ.MinDiff > table->uiSubDiffThreshold) {
                cout << "Miss mutated read" << i << endl;
                pass = false;
                // PRINT_MASKED_REPERAT_FLAG(table, i);
            } else if (alignmentsQ.MinDiff < table->uiSubDiffThreshold) {
                // cout << "Found better alignment!" << endl;
            }

            // (3) Check reversed compliment reads
            if (table->bMapReadInColors) {
                table->pgenomeNTInBits->fragACGTKmerInBits(slideWindows, i, table->uiRead_Length);
                CReadInBits rcRead = reverseCompliment(table->uiRead_Length, slideWindows);
                CReadInBits rcReadInColors = bases2Colors(rcRead);

                rcReadInColors.decode(read);
                myStrCpy(originalRead, read, MAX_READ_LENGTH);
                introduceMutation(read, table->chosenSeedId);
                CReadInBits rcMreadInBits(read);
                alignmentsQ.clearHits();
                table->queryReadColors(rcMreadInBits, alignmentsQ, true, false);
            } else {
                reverseComplementKmer(read);
                CReadInBits rcMreadInBits(read);
                alignmentsQ.MinDiff = table->uiRead_Length; // SET A LARGE NUMBER
                table->queryReadBases(rcMreadInBits, alignmentsQ, true, false);
            }

            if (alignmentsQ.MinDiff > table->uiSubDiffThreshold) {
                cout << "Miss reversed mutated read" << i << endl;
                pass = false;
                // PRINT_MASKED_REPERAT_FLAG(table, i);
            } else if (alignmentsQ.MinDiff < table->uiSubDiffThreshold) {
                // cout << "Found better alignment!" << endl;
            }
        }
    }
    return(pass);
}

